Skip to main content

Reordering a Tensor using Einsum with Sparse Partials

This is an example of how to properly use the einsum function to reorder a tensor while using sparse partial derivatives.

from csdl_om import Simulatorimport numpy as npfrom csdl import Modelimport csdl

class ExampleReorderTensorSparse(Model):
    def define(self):
        # Shape of Tensor        shape3 = (2, 4, 3)        c = np.arange(24).reshape(shape3)
        # Declaring tensor        tens = self.declare_variable('c', val=c)
        self.register_output(            'einsum_reorder2_sparse_derivs',            csdl.einsum_new_api(tens,                                operation=[(33, 66, 99), (99, 66, 33)],                                partial_format='sparse'))

sim = Simulator(ExampleReorderTensorSparse())sim.run()
print('c', sim['c'].shape)print(sim['c'])print('einsum_reorder2_sparse_derivs',      sim['einsum_reorder2_sparse_derivs'].shape)print(sim['einsum_reorder2_sparse_derivs'])
[[[ 0.  1.  2.]  [ 3.  4.  5.]  [ 6.  7.  8.]  [ 9. 10. 11.]]
 [[12. 13. 14.]  [15. 16. 17.]  [18. 19. 20.]  [21. 22. 23.]]]einsum_reorder2_sparse_derivs (3, 4, 2)[[[ 0. 12.]  [ 3. 15.]  [ 6. 18.]  [ 9. 21.]]
 [[ 1. 13.]  [ 4. 16.]  [ 7. 19.]  [10. 22.]]
 [[ 2. 14.]  [ 5. 17.]  [ 8. 20.]  [11. 23.]]]